import pandas as pd
import os
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

# 对于速度进行处理，可以先看一些结果图。简而言之就是：对于每个模式的速度分布选择直方图和KDE密度图进行展示
# 展示方式可以是整体，可以是分类（核心的两大类：运输航道和沿岸航道）下的SOG速度

def process_csj(file_path):
    # CSJ处理逻辑
    columns_csj = ['MMSI', 'Timestamp', 'TimeDelta', 'LATITUDE', 'LONGITUDE', 'SOG', 'COG']
    df = pd.read_csv(file_path, sep=' ', names=columns_csj,
                     usecols=['MMSI', 'Timestamp', 'TimeDelta', 'LATITUDE', 'LONGITUDE', 'SOG', 'COG'])
    df['Timestamp'] = pd.to_datetime(df['Timestamp'], format='%Y%m%d%H%M%S')
    df = df[['MMSI', 'Timestamp', 'LONGITUDE', 'LATITUDE', 'SOG', 'COG']]
    return df


def process_zs(file_path, mmsi, subdir):
    # ZS处理逻辑
    df = pd.read_csv(file_path, usecols=['Date_Time', 'SOG', 'COG', 'LONGITUDE', 'LATITUDE'])
    if subdir == "2018-4-24":
        df['Date_Time'] += 86400  # 加上一天的秒数
    df.rename(columns={'Date_Time': 'Timestamp'}, inplace=True)
    df['Timestamp'] = pd.to_datetime(df['Timestamp'], unit='s', origin=pd.Timestamp('2018-04-23'))
    df['MMSI'] = mmsi
    df = df[['MMSI', 'Timestamp', 'LONGITUDE', 'LATITUDE', 'SOG', 'COG']]
    return df


def process_cfd(file_path, mmsi, subdir):
    # 提取日期中的天数，并转换为对应的秒数（距离2018-06-01的天数乘以每天的秒数）
    day = int(subdir.split('-')[-1])
    seconds_from_start = (day - 1) * 86400
    columns_cfd = ['Timestamp', 'TimeDelta', 'LONGITUDE', 'LATITUDE', 'SOG', 'COG', 'WhateverOG Idontknow']
    df = pd.read_csv(file_path, names=columns_cfd, usecols=['TimeDelta', 'LONGITUDE', 'LATITUDE', 'SOG', 'COG'])
    df['TimeDelta'] += seconds_from_start
    df['Timestamp'] = pd.to_datetime(df['TimeDelta'], unit='s', origin=pd.Timestamp('2018-06-01'))
    df['MMSI'] = mmsi
    df = df[['MMSI', 'Timestamp', 'LONGITUDE', 'LATITUDE', 'SOG', 'COG']]
    return df


def grid_average_sog(cluster_data, grid_size=(500, 500)):
    min_lat, max_lat = cluster_data['LATITUDE'].min(), cluster_data['LATITUDE'].max()
    min_lon, max_lon = cluster_data['LONGITUDE'].min(), cluster_data['LONGITUDE'].max()

    lat_bins = np.linspace(min_lat, max_lat, grid_size[0] + 1)
    lon_bins = np.linspace(min_lon, max_lon, grid_size[1] + 1)

    grid = np.zeros(grid_size)
    counts = np.zeros(grid_size)

    for _, row in cluster_data.iterrows():
        lat_idx = np.digitize(row['LATITUDE'], lat_bins) - 1
        lon_idx = np.digitize(row['LONGITUDE'], lon_bins) - 1
        if 0 <= lat_idx < grid_size[0] and 0 <= lon_idx < grid_size[1]:
            grid[lat_idx, lon_idx] += row['SOG']
            counts[lat_idx, lon_idx] += 1

    avg_sog_grid = np.divide(grid, counts, out=np.zeros_like(grid), where=counts != 0)
    return avg_sog_grid, min_lat, max_lat, min_lon, max_lon


def visualize_sog_grid(avg_sog_grid, output_path, min_lat, max_lat, min_lon, max_lon):
    plt.imshow(avg_sog_grid, cmap='plasma', interpolation='nearest', extent=[min_lon, max_lon, min_lat, max_lat])
    plt.colorbar(label='Average SOG')
    plt.title('Average SOG by Grid')
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    plt.gca().invert_yaxis()  # 确保经纬度从下到上增大
    plt.savefig(output_path)
    plt.close()


def plot_sog_histogram(cluster_data, cluster_number, output_directory):
    plt.figure(figsize=(10, 6))
    plt.hist(cluster_data['SOG'], bins=30, color='blue', edgecolor='black')
    plt.title(f'SOG Distribution for Cluster {cluster_number}')
    plt.xlabel('SOG (knots)')
    plt.ylabel('Frequency')
    plt.grid(True)
    plt.savefig(os.path.join(output_directory, f'cluster_{cluster_number}_sog_histogram.png'))
    plt.close()


def plot_sog_density(cluster_data, cluster_number, output_directory):
    plt.figure(figsize=(10, 6))
    sns.kdeplot(cluster_data['SOG'], shade=True, color='blue')
    plt.title(f'SOG Density for Cluster {cluster_number}')
    plt.xlabel('SOG (knots)')
    plt.ylabel('Density')
    plt.grid(True)
    plt.savefig(os.path.join(output_directory, f'cluster_{cluster_number}_sog_density.png'))
    plt.close()


def plot_combined_sog_density(mmsi_cluster_df, all_trajectories, output_directory):
    clusters = mmsi_cluster_df['Cluster'].unique()
    plt.figure(figsize=(10, 6))

    colors = ['blue', 'green', 'red', 'cyan', 'magenta', 'yellow', 'black', 'orange']

    for cluster, color in zip(clusters, colors):
        cluster_mmsi = mmsi_cluster_df[mmsi_cluster_df['Cluster'] == cluster]['MMSI']
        cluster_data = pd.concat([all_trajectories[mmsi] for mmsi in cluster_mmsi if mmsi in all_trajectories])

        if not cluster_data.empty:
            sns.kdeplot(cluster_data['SOG'], shade=True, color=color, label=f'Cluster {cluster}')

    plt.title('Combined SOG Density for All Clusters')
    plt.xlabel('SOG (knots)')
    plt.ylabel('Density')
    plt.legend()
    plt.grid(True)
    plt.savefig(os.path.join(output_directory, 'combined_sog_density.png'))
    plt.close()

def plot_combined_sog_density(mmsi_cluster_df, all_trajectories, output_directory):
    clusters = [1, 4, 5, 6]
    plt.figure(figsize=(10, 6))

    colors = ['blue', 'green', 'red', 'cyan']

    for cluster, color in zip(clusters, colors):
        cluster_mmsi = mmsi_cluster_df[mmsi_cluster_df['Cluster'] == cluster]['MMSI']
        cluster_data = pd.concat([all_trajectories[mmsi] for mmsi in cluster_mmsi if mmsi in all_trajectories])

        if not cluster_data.empty:
            sns.kdeplot(cluster_data['SOG'], shade=True, color=color, label=f'Cluster {cluster}')

    plt.title('Combined SOG Density for Clusters 1, 4, 5, 6')
    plt.xlabel('SOG (knots)')
    plt.ylabel('Density')
    plt.legend()
    plt.grid(True)
    plt.savefig(os.path.join(output_directory, 'combined_sog_density_1_4_5_6.png'))
    plt.close()

def plot_combined_sog_density_rest(mmsi_cluster_df, all_trajectories, output_directory):
    clusters = [0, 2, 3, 7]
    plt.figure(figsize=(10, 6))

    colors = ['magenta', 'yellow', 'black', 'orange']

    for cluster, color in zip(clusters, colors):
        cluster_mmsi = mmsi_cluster_df[mmsi_cluster_df['Cluster'] == cluster]['MMSI']
        cluster_data = pd.concat([all_trajectories[mmsi] for mmsi in cluster_mmsi if mmsi in all_trajectories])

        if not cluster_data.empty:
            sns.kdeplot(cluster_data['SOG'], shade=True, color=color, label=f'Cluster {cluster}')

    plt.title('Combined SOG Density for Clusters 0, 2, 3, 7')
    plt.xlabel('SOG (knots)')
    plt.ylabel('Density')
    plt.legend()
    plt.grid(True)
    plt.savefig(os.path.join(output_directory, 'combined_sog_density_0_2_3_7.png'))
    plt.close()

def plot_combined_sog_histograms(mmsi_cluster_df, all_trajectories, output_directory):
    clusters = mmsi_cluster_df['Cluster'].unique()
    fig, axes = plt.subplots(2, 4, figsize=(20, 10))
    axes = axes.flatten()

    for ax, cluster in zip(axes, clusters):
        cluster_mmsi = mmsi_cluster_df[mmsi_cluster_df['Cluster'] == cluster]['MMSI']
        cluster_data = pd.concat([all_trajectories[mmsi] for mmsi in cluster_mmsi if mmsi in all_trajectories])

        if not cluster_data.empty:
            ax.hist(cluster_data['SOG'], bins=30, color='#CCECFC', edgecolor='#E4B9F1')
            ax.set_title(f'Cluster {cluster}')
            ax.set_xlabel('SOG (knots)')
            ax.set_ylabel('Frequency')
            ax.grid(True)

    # Adjust layout
    plt.tight_layout()
    plt.savefig(os.path.join(output_directory, 'combined_sog_histograms.png'))
    plt.close()


def process_clusters(mmsi_cluster_df, all_trajectories, output_directory, grid_size=(100, 100)):
    clusters = mmsi_cluster_df['Cluster'].unique()

    for cluster in clusters:
        cluster_mmsi = mmsi_cluster_df[mmsi_cluster_df['Cluster'] == cluster]['MMSI']
        cluster_data = pd.concat([all_trajectories[mmsi] for mmsi in cluster_mmsi if mmsi in all_trajectories])

        if not cluster_data.empty:
            avg_sog_grid, min_lat, max_lat, min_lon, max_lon = grid_average_sog(cluster_data, grid_size)
            output_path = os.path.join(output_directory, f'cluster_{cluster}_sog.png')
            visualize_sog_grid(avg_sog_grid, output_path, min_lat, max_lat, min_lon, max_lon)

            # 绘制SOG直方图和核密度估计图
            plot_sog_histogram(cluster_data, cluster, output_directory)
            plot_sog_density(cluster_data, cluster, output_directory)
# 绘制合并的核密度估计图和直方图
    plot_combined_sog_density(mmsi_cluster_df, all_trajectories, output_directory)
    plot_combined_sog_density_rest(mmsi_cluster_df, all_trajectories, output_directory)
    plot_combined_sog_histograms(mmsi_cluster_df, all_trajectories, output_directory)

def main():
    mmsi_cluster_df = pd.read_csv("Data/DTW/CSJ/New_Human_Trajectories/cluster_mmsi_pairs.csv")

    dirs = {
        './Data/CSJ': process_csj,
        # './Data/ZS': process_zs,
        # './Data/CFD': process_cfd
    }
    all_trajectories = {}

    for root_dir, process_func in dirs.items():
        if root_dir.endswith('CSJ'):  # CSJ逻辑
            for file_name in os.listdir(root_dir):
                if file_name.endswith('.txt'):
                    file_path = os.path.join(root_dir, file_name)
                    df = process_csj(file_path)
                    if not df.empty:
                        mmsi = df.iloc[0]['MMSI']
                        if mmsi not in all_trajectories:
                            all_trajectories[mmsi] = []
                        all_trajectories[mmsi].append(df)

        elif root_dir.endswith('CFD'):
            for subdir in os.listdir(root_dir):
                subdir_path = os.path.join(root_dir, subdir)
                if os.path.isdir(subdir_path):
                    for file_name in os.listdir(subdir_path):
                        if file_name.endswith('.csv'):
                            file_path = os.path.join(subdir_path, file_name)
                            mmsi = file_name.split('.')[0]
                            df = process_cfd(file_path, mmsi, subdir)
                            if mmsi not in all_trajectories:
                                all_trajectories[mmsi] = []
                            all_trajectories[mmsi].append(df)

        else:  # ZS逻辑，包含子文件夹
            for subdir in os.listdir(root_dir):
                subdir_path = os.path.join(root_dir, subdir)
                if os.path.isdir(subdir_path):
                    for file_name in os.listdir(subdir_path):
                        if file_name.endswith('.csv'):
                            file_path = os.path.join(subdir_path, file_name)
                            mmsi = file_name.split('.')[0]
                            df = process_func(file_path, mmsi, subdir)
                            if mmsi not in all_trajectories:
                                all_trajectories[mmsi] = []
                            all_trajectories[mmsi].append(df)

    for mmsi, dfs in all_trajectories.items():
        all_trajectories[mmsi] = pd.concat(dfs).reset_index(drop=True).sort_values(by='Timestamp')

    output_directory = "Data/Groups/CSJ/New_Human_Trajectories/MMSI_Cluster_Groups/SOG"
    os.makedirs(output_directory, exist_ok=True)

    process_clusters(mmsi_cluster_df, all_trajectories, output_directory)


if __name__ == "__main__":
    main()
